package jsu.blogger.util;



import com.hankcs.hanlp.HanLP;
import com.hankcs.hanlp.seg.common.Term;

import java.util.*;
import java.util.stream.Collectors;

public class KeywordExtractor {
    // 停用词列表，添加更多常见的停用词
    private static final Set<String> STOP_WORDS = new HashSet<>(Arrays.asList(
            "一个", "的", "了", "是", "在", "和", "有", "也", "就", "人", "都", "而", "与", "这", "它", "为", "与", "对", "对比", "可以", "如何", "该", "我们", "你", "我", "他", "她", "自己"
            // 可以根据需求继续添加更多停用词
    ));

    // 停用词性过滤：过滤掉非名词（例如助词、介词、连词等）
    // 只保留名词相关的词性
    private static final Set<String> USEFUL_POS = new HashSet<>(Arrays.asList(
            "n", "nr", "ns", "nt", "nz", "vn", "nx"  // 增加动名词（vn）和外来名词（nx）
            // "v", "a", "vn", "ad" 被去除，避免包含动词、形容词等
    ));

    // 方法：根据给定文本提取前 topN 个关键词
    public static List<String> extractKeywords(String text, int topN) {
        // 1. 使用 HanLP 进行分词，分词后的结果是包含词和词性的信息
        List<Term> terms = HanLP.segment(text);

        // 2. 计算词频（TF）
        Map<String, Integer> termFrequency = new HashMap<>();
        for (Term term : terms) {
            String word = term.word;
            String pos = term.nature.toString();

            // 过滤掉停用词、标点符号和不相关词性
            if (!word.trim().isEmpty() && !STOP_WORDS.contains(word) && !isPunctuation(word) && USEFUL_POS.contains(pos)) {
                termFrequency.put(word, termFrequency.getOrDefault(word, 0) + 1);
            }
        }

        // 3. 计算逆文档频率（IDF）
        // 由于这里只有一个文本，所以 IDF 的计算简化为 log(1 + total_terms / term_frequency)
        int totalTerms = terms.size();
        Map<String, Double> idfMap = new HashMap<>();
        for (String term : termFrequency.keySet()) {
            // 逆文档频率（IDF）
            double idf = Math.log(1.0 + totalTerms / (double) termFrequency.get(term));
            idfMap.put(term, idf);
        }

        // 4. 计算TF-IDF（TF * IDF）
        Map<String, Double> tfIdfMap = new HashMap<>();
        for (Map.Entry<String, Integer> entry : termFrequency.entrySet()) {
            String word = entry.getKey();
            int tf = entry.getValue();
            double idf = idfMap.get(word);
            tfIdfMap.put(word, tf * idf);
        }

        // 5. 按照TF-IDF排序
        return tfIdfMap.entrySet().stream()
                .sorted((e1, e2) -> e2.getValue().compareTo(e1.getValue()))  // 按照TF-IDF值降序排序
                .limit(topN)  // 取前topN个关键词
                .map(Map.Entry::getKey)
                .collect(Collectors.toList());
    }

    // 判断是否为标点符号
    private static boolean isPunctuation(String word) {
        return word.matches("\\p{Punct}");
    }
}
